package result_reranking.entity_rank;

import common.dao.elastic.ElasticSearchDao;
import common.helper.json.JsonFormatter;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.MatchQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.SimpleQueryStringBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import query_understanding.query_understand.tree_progress.DocumentQueryProgress;
import query_understanding.query_understand.tree_selector.QueryTreeSelector;
import semantic_matching.common.schema.MatchedDocumentTree;
import semantic_matching.common.schema.MatchedEntites;
import semantic_matching.common.schema.MatchedEntityTree;
import semantic_matching.doc_search.doc_tree.DocumentQueryMatcher;
import textprocessing.datasources.enwikipedia.IndexerEnWikiNews;

import java.util.ArrayList;
import java.util.List;

/**
 * Created by julianzliu on 5/2/2017.
 */
public class DocumentRanking {

    static final Logger logger = LoggerFactory.getLogger( DocumentRanking.class );


    private ElasticSearchDao esDao;
    private IndexerEnWikiNews indexer;


    private QueryTreeSelector selector;
    private DocumentQueryMatcher matcher;



    public DocumentRanking( ) {
        this.indexer = new IndexerEnWikiNews();
        this.esDao = this.indexer.getEsDao();

        this.selector = new QueryTreeSelector();
        this.matcher = new DocumentQueryMatcher();
    }




    private  DocumentQueryProgress genQueryTree( String query ){
        DocumentQueryProgress progress = this.selector.genDocQueryTree( query );
        return progress;
    }

    private MatchedDocumentTree genMatchedTree( DocumentQueryProgress progress){
        MatchedDocumentTree matchedDocumentTree = matcher.solveDocumentQueryProgress( progress );
        return matchedDocumentTree;
    }

    private String parseEntityToString(String entity){
        String[] tokens = entity.split("<|>|_");
        String ret = "";
        for(String token: tokens) ret += " " + token;
        return ret;
    }

    private BoolQueryBuilder genAnswerEntityQuery(MatchedDocumentTree matchedDocumentTree ){
        if(matchedDocumentTree.matchedEntityTreeList.size() == 0)
            return null;
        Boolean addOneSubQuery = false;
        BoolQueryBuilder docEntityQuery = QueryBuilders.boolQuery();
        for(MatchedEntityTree entityTree: matchedDocumentTree.matchedEntityTreeList){
            BoolQueryBuilder entityQuery = QueryBuilders.boolQuery();
            Boolean atLeastOne = false;
            MatchedEntites rootEntites = entityTree.rootEntities;
            for(String entity: rootEntites.entities){
                atLeastOne = true;
                //System.out.println( "[rootEntites] tree = " + entityTree.hashCode() + " , entity = " +  entity );
                MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery("entities", entity);
                entityQuery.should( matchQueryBuilder );

                MatchQueryBuilder titleQuery = QueryBuilders.matchQuery("title", parseEntityToString( entity ) );
                entityQuery.should( titleQuery );

                MatchQueryBuilder textQuery = QueryBuilders.matchQuery("text", parseEntityToString( entity ) );
                entityQuery.should( textQuery );

            }
            if(atLeastOne){
                docEntityQuery.must( entityQuery );
                addOneSubQuery = true;
            }
        }
        if(addOneSubQuery)
            return docEntityQuery;
        else
            return  null;
    }



    public String rankResult( String text , Boolean is_dsl) {
        logger.info("[rankResult] search request = {}\n\n", text);

        DocumentQueryProgress progress = genQueryTree( text );
        MatchedDocumentTree matchedDocumentTree = genMatchedTree( progress );



        BoolQueryBuilder query = QueryBuilders.boolQuery();

        /*
        必须匹配的子查询，召回 包含原始查询结果  或者  答案实体  的文档
         */
        BoolQueryBuilder must_query = QueryBuilders.boolQuery();
        BoolQueryBuilder should_query = QueryBuilders.boolQuery();
        query.must( must_query );
        query.should(should_query);


        /*
        原始查询的全文检索
         */
        MatchQueryBuilder fulltextQuery = QueryBuilders.matchQuery("text",text);
        must_query.should( fulltextQuery.boost(RankParameters.fulltextSearchBoost) );


        /*
        包含结果实体的文档
         */
        //List<String> answerEntities = new ArrayList<>();
        //BoolQueryBuilder answerEntityQuery = genEntitiesQuery( answerEntities );
        BoolQueryBuilder answerEntityQuery = genAnswerEntityQuery( matchedDocumentTree );
        if(answerEntityQuery != null)
            must_query.should( answerEntityQuery.boost( RankParameters.answerEntityBoost ) );



        /*
        包含主题实体的文档
         */
        List<String> topicEntities = new ArrayList<>();
        BoolQueryBuilder topicEntityQuery = genEntitiesQuery( topicEntities );
        if(topicEntityQuery != null)
            should_query.should( topicEntityQuery );


        if(is_dsl != null && is_dsl){
            SearchRequestBuilder request = this.esDao.genSearchWithApiRequest(this.indexer.es_doc_index, this.indexer.es_doc_type, query, 0, 10, null, null);

            return request.toString();
        }


        SearchResponse response = this.esDao.searchWithApi(this.indexer.es_doc_index, this.indexer.es_doc_type, query, 0, 10, null, null);

        return JsonFormatter.formatSearchResponse(response);
    }


    private BoolQueryBuilder genEntitiesQuery( List<String> entities ){
        if(entities.size() == 0) return null;
        BoolQueryBuilder entitiesQuery = QueryBuilders.boolQuery();
        for(String entity: entities){
            MatchQueryBuilder entityQuery = QueryBuilders.matchQuery("entities", entity);
            entitiesQuery.should( entityQuery );
        }
        return entitiesQuery;
    }


    public String simpleQuery(String text){
        SimpleQueryStringBuilder string_query = QueryBuilders.
                simpleQueryStringQuery( text )
                .field( "text" );
        SearchResponse response = this.esDao.searchWithApi(this.indexer.es_doc_index, this.indexer.es_doc_type, string_query, 0, 10, null, null);

        return JsonFormatter.formatSearchResponse(response);
    }



    /*********************************************************************
     *
     * @param args
     */
    public static void main(String[] args){

        String text = "Obama";
        if(args.length > 0){
            text = "";
            for(String arg: args)
                text += " " + arg;
        }
        text += " .";

        DocumentRanking documentRanking = new DocumentRanking();



        String result = documentRanking.rankResult( text, false  );

        System.out.println("\n\n[Ranking Result]--------------------------\n\n");

        System.out.println( result );


        /*
        print simple query DSL
         */
        documentRanking.simpleQuery( text );

    }

}
